import os
from datetime import datetime
import torch
import swanlab 
import yaml
from .colorful import Color

import logging
# 配置一个简单的日志器，以便 LogSystem 中的 logger.info 可以工作
logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
logger = logging.getLogger(__name__)


class LogSystem:
    def __init__(self, log_dir:str = "./logs", project_name:str = 'tmp', experiment_name:str = 'tmp', local_rank:int = 0):
        self.local_rank = local_rank
        if self.local_rank != 0:
            return

        self._steps = {}
        self._base_log_dir = log_dir  # 存储基础日志目录
        self._project_name = project_name # 存储项目/实验组名称
        self._run_timestamp = os.path.join(experiment_name, datetime.now().strftime('%Y%m%d_%H%M%S')) # 当前运行的时间戳

        # _log_dir 现在表示我们为本次运行预期的“逻辑”日志目录，
        # 对于 Tensorboard，它就是实际写入的目录。
        # 对于 Swanlab，它将是 swanlab 内部创建目录的父路径之一。
        self._log_dir = os.path.join(self._base_log_dir, self._project_name, self._run_timestamp)

        self._init_logger()

        logger.info(f"Log directory is initialized at {Color.Yellow}{self._log_dir}{Color.End}")

    def _get_log_dir(self):
        """
        返回本次运行的特定日志目录。
        对于 Tensorboard，SummaryWriter 会直接写入此目录。
        对于 Swanlab，此目录是其内部创建运行目录的父路径之一。
        """
        if not os.path.exists(self._log_dir):
            os.makedirs(self._log_dir, exist_ok=True)
        return self._log_dir

    def log(self, message:dict):
        if self.local_rank != 0:
            return
        assert isinstance(message, dict), f"Message must be a dict, but got {type(message)}"
        for key, value in message.items():
            if isinstance(value, torch.Tensor):
                value = value.detach().cpu()
                if torch.any(torch.isnan(value)):
                    raise ValueError(f"Value {key} is NaN")
                value = value.item()
            step = self._steps.get(key, 0)
            self._steps[key] = step + 1
            self._flush_impl(step, key, value)

    def _flush_impl(self, step, key, value):
        raise NotImplementedError

    def _init_logger(self):
        raise NotImplementedError
    
    def config(self, config:dict):
        raise NotImplementedError

# 从 torch.utils.tensorboard 导入 SummaryWriter (保持原样)
from torch.utils.tensorboard import SummaryWriter

class TensorboardLogger(LogSystem):
    def __init__(self, log_dir: str = "./logs", project_name:str = 'tmp', experiment_name:str = 'tmp', local_rank:int=0):
        super().__init__(log_dir, project_name, experiment_name, local_rank)

    def _init_logger(self):
        # Tensorboard SummaryWriter 直接写入 _get_log_dir 返回的特定运行目录
        self._logger = SummaryWriter(self._get_log_dir())

    def _flush_impl(self, step, key, value):
        self._logger.add_scalar(key, value, step)

    def config(self, config:dict):
        yaml.dump(config, open(os.path.join(self._get_log_dir(), 'configs.yaml'), 'w'))

from dataclasses import dataclass, asdict
class SwanlabLogger(TensorboardLogger):
    def __init__(self, log_dir: str = "./logs", project_name:str = 'tmp', experiment_name:str = 'tmp', local_rank:int=0):
        super().__init__(log_dir, project_name, experiment_name, local_rank)
        # 如果 local_rank 不为 0，父类构造函数会返回，这里也无需进一步初始化
        if self.local_rank != 0:
            return

    def _init_logger(self):
        """
        初始化 Swanlab 日志系统。
        swanlab.init 的 logdir 参数是项目日志的根目录。
        project 参数是项目名称。
        experiment_name 参数是本次运行的名称。
        swanlab 会在 logdir/project/experiment_name 结构下创建其运行目录。
        """
        super()._init_logger()
        
        # 使用 _base_log_dir 作为 swanlab 的 logdir
        # 使用 _project_name 作为 swanlab 的 project
        # 使用 _run_timestamp 作为本次 swanlab 运行的 experiment_name
        # 这样，swanlab 会在例如 "./logs/my_project/20230101_123456" 下创建其内部的运行目录
        swanlab.init(
            logdir=self._base_log_dir,
            project=self._project_name,
            experiment_name=self._run_timestamp,
            # 可以添加其他 swanlab.init 参数，例如 description, config, tags 等
            # 例如：description=f"Experiment run at {self._run_timestamp}"
        )
        # Swanlab 不像 Tensorboard 那样需要一个 self._logger 对象来调用 add_scalar，
        # 而是直接通过 swanlab.log() 函数进行日志记录。

    def _flush_impl(self, step, key, value):
        """
        将数据记录到 Swanlab。
        swanlab.log 接受一个字典，其中包含要记录的指标，并可选地接受一个 step 参数。
        """
        # swanlab.log 期望一个字典，所以我们构建 {key: value}
        super()._flush_impl(step, key, value)
        swanlab.log({key: value}, step=step)
     
    def config(self, config:dict):
        super().config(config)
        for key, value in config.items():
            try:
                value = asdict(value)
            except:
                pass
            swanlab.config.set(key, value)
